from __future__ import annotations

from typing import TYPE_CHECKING

import cv2
import gymnasium as gym
import numpy as np
import torch

if TYPE_CHECKING:
    from uav_2d.uav_env import Uav2DEnv


class RasterWrapper(gym.Wrapper):
    def __init__(
        self,
        env: gym.Env,
        state_size: tuple[int, int] = (384, 384),
        state_downsize: tuple[int, int] = (128, 128),
        use_sgcnn: bool = True,
    ):
        super().__init__(env)
        self.raw_env: "Uav2DEnv" = self.env.unwrapped
        self.state_size = state_size
        self.state_downsize = state_downsize
        self.use_sgcnn = use_sgcnn
        obs_shape = (4, *self.state_downsize)

        if use_sgcnn:
            self.max_pool = torch.nn.MaxPool2d(kernel_size=(2, 2), stride=2)
            obs_shape = (16, *(ds // 8 for ds in self.state_downsize))

        for i in range(self.raw_env.num_agents):
            # Directly modify the inner observation space
            self.raw_env.observation_space[f"uav_{i}"].spaces["raster"] = (
                gym.spaces.Box(low=0.0, high=1.0, shape=obs_shape, dtype=np.float32)
            )

        # Global map, shared by all agents
        self.map_obstacle = np.zeros(
            (self.raw_env.dimensions[1], self.raw_env.dimensions[0]), dtype=np.float32
        )
        self.last_map_opponent = np.zeros(
            (self.raw_env.dimensions[1], self.raw_env.dimensions[0]), dtype=np.float32
        )

        # Multi-agent independent maps
        self.map_distances: dict[str, np.ndarray] = {}
        self.map_trajectories: dict[str, np.ndarray] = {}
        self.last_positions: dict[str, tuple[int, int]] = {}

    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        for agent_id, agent_obs in obs.items():
            agent_uav = self.raw_env.agents[agent_id]
            current_pos = (
                int(np.round(agent_uav.x)),
                int(np.round(agent_uav.y)),
            )
            cv2.line(
                self.map_trajectories[agent_id],
                self.last_positions[agent_id],
                current_pos,
                color=(1.0,),
            )
            self.last_positions[agent_id] = current_pos
            obs[agent_id]["raster"] = self.get_raster(agent_id)
        return obs, reward, terminated, truncated, info

    def get_raster(self, agent_id: str) -> np.ndarray:
        map_opponent = np.zeros(
            (self.raw_env.dimensions[1], self.raw_env.dimensions[0]), dtype=np.float32
        )
        for d_obstacle in self.raw_env.d_obstacles:
            pts = np.array(
                [
                    [
                        *(
                            (
                                np.round(pt[0]).astype(np.int32),
                                np.round(pt[1]).astype(np.int32),
                            )
                            for pt in d_obstacle.get_hull()
                        )
                    ]
                ],
                dtype=np.int32,
            )
            cv2.fillPoly(map_opponent, [pts], color=(1.0,))

        self.last_map_opponent = 0.5 * self.last_map_opponent + 0.5 * map_opponent
        self.last_map_opponent = np.where(
            np.abs(self.last_map_opponent) < 0.5 ** (4 - 1),
            0.0,
            self.last_map_opponent,
        )

        for structure in self.raw_env.obstacles.get_occur_structures():
            up_left = structure.pts[0]
            down_right = structure.pts[2]
            self.map_obstacle[up_left.y : down_right.y, up_left.x : down_right.x] = 1.0

        raster_obs = np.stack(
            (
                self.map_distances[agent_id],
                self.map_obstacle,
                self.map_trajectories[agent_id],
                self.last_map_opponent,
            ),
            axis=-1,
        )
        diag_r = self.state_size[0] / 2 * np.sqrt(2)
        diag_r_int = int(np.ceil(diag_r))
        raster_obs = cv2.copyMakeBorder(
            raster_obs,
            diag_r_int,
            diag_r_int,
            diag_r_int,
            diag_r_int,
            cv2.BORDER_CONSTANT,
            value=np.array((0.0, 1.0, 0.0, 0.0)),
        )
        agent_uav = self.raw_env.agents[agent_id]
        x_t = int(np.round(agent_uav.x))
        y_t = int(np.round(agent_uav.y))
        x_t = max(0, min(x_t, self.raw_env.dimensions[0] - 1))
        y_t = max(0, min(y_t, self.raw_env.dimensions[1] - 1))
        leftmost = int(round(y_t))
        rightmost = int(round(y_t + 2 * diag_r_int))
        upmost = int(round(x_t))
        bottommost = int(round(x_t + 2 * diag_r_int))
        obs_cropped = raster_obs[leftmost:rightmost, upmost:bottommost, :]

        rotation_mat = cv2.getRotationMatrix2D(
            (diag_r, diag_r), -agent_uav.direction, 1.0
        )
        dst_size = 2 * diag_r_int
        delta = int(diag_r_int - self.state_size[0] / 2)
        obs_rotated = cv2.warpAffine(obs_cropped, rotation_mat, (dst_size, dst_size))
        obs_rotated = obs_rotated[
            delta : delta + self.state_size[0], delta : delta + self.state_size[0], :
        ]

        obs_rotated_resize = cv2.resize(obs_rotated, self.state_downsize)
        obs_final = obs_rotated_resize.transpose(2, 0, 1).astype(np.float32)
        if self.use_sgcnn:
            obs_final = self.sgcnn(obs_final)
        return obs_final

    def sgcnn(self, obs):
        sgcnn_size = 16
        obs_ = obs.copy()
        obs_list = []
        center_size = self.state_downsize[0] // 2
        with torch.no_grad():
            for _ in range(4):
                obs_list.append(
                    obs_[
                        :,
                        (center_size - sgcnn_size // 2) : (
                            center_size + sgcnn_size // 2
                        ),
                        (center_size - sgcnn_size // 2) : (
                            center_size + sgcnn_size // 2
                        ),
                    ]
                )
                obs_ = self.max_pool(torch.from_numpy(obs_)).numpy()
                center_size //= 2
        return np.concatenate(obs_list, axis=0, dtype=np.float32)

    def reset(self, *, seed=None, options=None):
        obs, info = self.env.reset(seed=seed, options=options)
        self.map_distances = {}
        for agent_id, uav in self.raw_env.agents.items():
            x_coords = np.arange(0, self.raw_env.dimensions[0])
            y_coords = np.arange(0, self.raw_env.dimensions[1])
            grid_x, grid_y = np.meshgrid(x_coords, y_coords)
            goal_x, goal_y = uav.goal_position
            map_apf = np.sqrt((grid_x - goal_x) ** 2 + (grid_y - goal_y) ** 2)
            self.map_distances[agent_id] = 0.998**map_apf

        self.map_obstacle = np.zeros(
            (self.raw_env.dimensions[1], self.raw_env.dimensions[0]), dtype=np.float32
        )
        self.last_map_opponent = np.zeros(
            (self.raw_env.dimensions[1], self.raw_env.dimensions[0]), dtype=np.float32
        )
        self.map_trajectories = {
            agent_id: np.zeros(
                (self.raw_env.dimensions[1], self.raw_env.dimensions[0]),
                dtype=np.float32,
            )
            for agent_id in obs.keys()
        }
        self.last_positions = {
            agent_id: (
                int(np.round(self.raw_env.agents[agent_id].x)),
                int(np.round(self.raw_env.agents[agent_id].y)),
            )
            for agent_id in obs.keys()
        }
        for structure in self.raw_env.obstacles.get_fixed_structures():
            up_left = structure.pts[0]
            down_right = structure.pts[2]
            self.map_obstacle[up_left.y : down_right.y, up_left.x : down_right.x] = 1.0

        for agent_id in obs.keys():
            obs[agent_id]["raster"] = self.get_raster(agent_id)
        return obs, info
